#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# JiuZero 2025/5/22

import argparse
import os
import sys
import config
from lib.core.log import logger

def int_list(value):
    try:
        values = list(map(int, value.split(',')))
        for v in values:
            if v not in (-1, 0, 1, 2, 3):
                logger.error(f"Invalid option {v} (allowed values: -1-3)")
                sys.exit(0)
        return values
    except ValueError:
        logger.error("The parameter is in the wrong format, please use comma-separated integers (e.g. 1, 2, 3)")
        sys.exit(0)

def str_list(value):
    try:
        values = list(map(str, value.split(',')))
        return values
    except ValueError:
        logger.error("The parameter is malformed, please use comma-separated characters")
        sys.exit(0)

def cmd_line_parser(argv=None):
    """
    This function parses the command line parameters and arguments
    """

    if not argv:
        argv = sys.argv
        
    _ = os.path.basename(argv[0])
    usage = "z0 [command options] [arguments...]"
    parser = argparse.ArgumentParser(prog='Z0SCAN', usage=usage)
    
    subparsers = parser.add_subparsers(
        dest='command',
        title='Commands', 
        metavar='', 
        help=''
        )
    version_parser = subparsers.add_parser('version', help='Show program\'s version number and exit')
    scan_parser = subparsers.add_parser('scan', help='Scan command')
    reverse_parser = subparsers.add_parser('reverse', help='Reverse command')
    list_parser = subparsers.add_parser('list', help='List of plugins and dicts')
    dbcmd_parser = subparsers.add_parser('dbcmd', help='Database manage console')

    ## z0 scan
    # Proxy options
    proxy = scan_parser.add_argument_group('Proxy', 'Passive Agent Mode Options')
    proxy.add_argument("-s", "--server-addr", dest="server_addr", help="Server addr format:(ip:port) ")
    # Target options
    target = scan_parser.add_argument_group('Target', "Options has to be provided to define the target(s)")
    target.add_argument("-u", "--url", dest="url", help="Target URL (e.g. \"http://www.site.com/vuln.php?id=1\")")
    target.add_argument("-c", "--crawl", dest="crawl", help="Crawl on URL and set depth for crawler", type=int, default=False)
    target.add_argument("-f", "--file", dest="url_file", help="Scan multiple targets given in a textual file")
    target.add_argument("-Rs", "--redis-server", dest="redis_server", help="Connect to redis, run as server (e.g. --redis-server password@host:port:db)")
    # Connecton options
    conn = scan_parser.add_argument_group('Connection', 'Server Connection Options')
    conn.add_argument("-R", "--reverse-client", dest="reverse_client", action="store_true", default=False, help="Connect to reverse server. (Default False)")
    conn.add_argument("-Rc", "--redis-client", dest="redis_client", help="Connect to redis, run as client (e.g. --redis-client password@host:port:db)")
    # Requests options
    request = scan_parser.add_argument_group("Request", "Network request options")
    request.add_argument("-p", "--proxy", dest="proxy", type=str, help="Use a proxy/proxies from file to connect to the target URL, Support http,https,socks5,socks4 & txt,json (e.g. http://127.0.0.1:8080 or proxy.txt)")
    request.add_argument("--timeout", dest="timeout", default=config.TIMEOUT, help="Seconds to wait before timeout connection (Default {})".format(config.TIMEOUT), type=int)
    request.add_argument("--retry", dest="retry", type=list, default=config.RETRY, help="Time out retrials times (Default {})".format(config.RETRY))
    request.add_argument("--random-agent", dest="random_agent", action="store_true", default=False, help="Use randomly selected HTTP User-Agent header value")
    # Output options
    output = scan_parser.add_argument_group("Output", "Output options")
    output.add_argument("--html", dest="html", help="When selected, the output will be output to the output directory by default, or you can specify", action='store_true')
    output.add_argument("--json", dest="json", help="The json file is generated by default in the output directory, you can change the path")
    # Optimization options
    optimization = scan_parser.add_argument_group("Optimization", "Optimization options")
    optimization.add_argument("-l", "--level", dest="level", type=int, choices=list(range(0, 4)), default=config.LEVEL, help="Different level use different payloads: 0-3 (Default {})".format(config.LEVEL))
    optimization.add_argument("-r", "--risk", dest="risk", type=int_list, default=config.RISK, help="Different risk use different kind of scanners: [0, 1, 2, 3] (Default {})".format(config.RISK))
    optimization.add_argument('-t', "--threads", dest="threads", type=int, default=config.THREADS, help="Threads of plugins tasks (Default {})".format(config.THREADS))
    optimization.add_argument("-cp", "--console-port", dest="console_port", type=int, help=f"Set port for console (Default {config.CONSOLE_PORT})", default=config.CONSOLE_PORT)
    optimization.add_argument('-pt', "--plugin-threads", dest="threads", type=int, default=config.PLUGIN_THREADS, help="Threads in plugins (Default {})".format(config.PLUGIN_THREADS))
    optimization.add_argument("-iw", '--ignore-waf', dest='ignore_waf', action="store_true", default=False, help="Ignore the WAF during detection")
    optimization.add_argument("-if", '--ignore-fingerprint', dest='ignore_fingerprint', action="store_true", default=False, help="Ignore fingerprint element scanning")
    optimization.add_argument("-sc", '--scan-cookie', dest='scan_cookie', action="store_true", default=False, help="Scan cookie during detection")
    optimization.add_argument('--disable', dest='disable', type=str_list, default=config.DISLOAD, help="Disload some scanners.")
    optimization.add_argument('--enable', dest='enable', type=str_list, default=[], help="Only load scanners")
    optimization.add_argument("--ipv6", dest="ipv6", action="store_true", help="Use IPV6 address. It will try for IPV6 at first，else IPV4.")
    optimization.add_argument('--redis-clean', dest='redis_clean', action="store_true", default=False, help="Clean Redis")
    optimization.add_argument('--short-out', dest='short_out', action="store_true", default=False, help="Output on cmd shortly like W13Scan.")
    optimization.add_argument("--debug", dest="debug", type=int, choices=list(range(1, 4)), help="Show programs's exception: 1-3")
    
    args = parser.parse_args()
    dd = args.__dict__

    if not (dd['command'] in ["version", "reverse", "dbcmd", "list", "scan", "update"]):
        errMsg = "An error command input (version, scan, reverse, dbcmd, list, console). "
        errMsg += "Use -h for help\n"
        parser.error(errMsg)

    if args.command == 'scan' and not any((dd.get("server_addr"), dd.get("url"), dd.get("url_file"), dd.get("redis_server"))):
        errMsg = "Missing a mandatory option (-s, --server-addr, -u, --url, -f, --file, -Rs, --redis_server). "
        errMsg += "Use -h for basic and -hh for advanced help\n"
        parser.error(errMsg)
    
    if args.command == 'scan' and not (dd.get("url") or dd.get("url_file")) and dd.get("crawl"):
        errMsg = "Crawler should be work with \"-u(--url)\" or \"-f(--file)\"."
        parser.error(errMsg)

    return args